from fastapi import APIRouter
from typing import Optional
import os

# FastAPI中流式返回数据给客户端，可以使用StreamingResponse
from fastapi.responses import StreamingResponse
from openai import OpenAI
from pydantic import BaseModel
from app.settings.settings import ailiyun_key

ai = APIRouter()


class ChatIn(BaseModel):
    content: str


# 生成数据流的异步函数
async def event_generator(completion):
    """
    根据提供的完成数据生成数据流。

    :param:
    - completion: 一个包含多个部分的完成数据，每个部分都有可能的内容。

    该函数会根据内容的长度或特定的结束标记来决定何时输出缓冲区的内容。
    """
    buffer = ""  # 初始化缓冲区
    min_length = 50  # 设置最小输出长度

    # 遍历完成数据的每个部分
    for chunk in completion:
        # 如果当前部分包含内容，则将其添加到缓冲区
        if chunk.choices[0].delta.content:
            buffer += chunk.choices[0].delta.content

            # 当缓冲区达到最小长度或遇到句子结束标记时输出
            if len(buffer) >= min_length or any(
                mark in buffer for mark in [".", "。", "!", "?", "！", "？", "\n"]
            ):
                yield f"{buffer}\n"  # 输出缓冲区内容并
                buffer = ""  # 重置缓冲区

    # 确保最后的缓冲区内容也被输出
    if buffer:
        yield f"{buffer}\n"


@ai.post("/chat")
async def chat(chat_in: ChatIn):
    client = OpenAI(
        api_key=ailiyun_key,  # 这个是密钥是我个人的，最终发布时抹去
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    completion = client.chat.completions.create(
        model="qwen-plus",  # 模型列表：https://help.aliyun.com/zh/model-studio/get
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": chat_in.content},
        ],
        stream=True,
    )
    # 返回数据流
    return StreamingResponse(
        event_generator(completion), media_type="text/event-stream"
    )


# 用于存储对话历史的字典，key是用户标识，value是消息列表
chat_history = {}


class ChatTestIn(BaseModel):
    content: str
    user_id: str  # 添加用户标识字段


# 测试对话，能够实现上下文对话和流式返回
@ai.post("/chat/test")
async def chat_test(chat_test_in: ChatTestIn):
    """
    接收聊天输入并返回模型的流式回答。

    :param:
    - chat_test_in: ChatTestIn 类型的对象，包含用户ID和聊天内容。

    :return:
    - StreamingResponse 类型，包含流式返回的聊天回复。
    """
    try:
        # 初始化OpenAI客户端，使用阿里云的API密钥和基础URL
        client = OpenAI(
            api_key=ailiyun_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        # 获取或初始化用户的对话历史
        if chat_test_in.user_id not in chat_history:
            chat_history[chat_test_in.user_id] = [
                {"role": "system", "content": "You are a helpful assistant."}
            ]

        # 添加用户的新消息到对话历史中
        chat_history[chat_test_in.user_id].append(
            {"role": "user", "content": chat_test_in.content}
        )

        # 创建聊天完成请求，使用指定模型和对话历史，开启流式返回
        completion = client.chat.completions.create(
            model="qwen-plus",
            messages=chat_history[chat_test_in.user_id],  # 使用完整的对话历史
            stream=True,
        )

        # 初始化助手的回复消息
        assistant_message = {"role": "assistant", "content": ""}

        # 定义一个异步生成器，用于处理和返回助手的流式回复
        async def wrapped_generator(completion):
            nonlocal assistant_message  # 使用闭包内的变量
            buffer = ""  # 缓存助手的回复内容
            min_length = 50  # 缓存内容的最小长度

            # 遍历模型的流式回复
            for chunk in completion:
                if chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    assistant_message["content"] += content  # 更新助手的回复内容
                    buffer += content  # 更新缓存内容

                    # 当缓存内容达到最小长度或遇到标点符号时，发送缓存内容
                    if len(buffer) >= min_length or any(
                        mark in buffer
                        for mark in [".", "。", "!", "?", "！", "？", "\n"]
                    ):
                        yield f"{buffer}\n"  # 发送缓存内容
                        buffer = ""  # 清空缓存

            # 发送剩余的缓存内容
            if buffer:
                yield f"{buffer}\n"

            # 将助手的完整回复添加到历史记录
            chat_history[chat_test_in.user_id].append(assistant_message)

            # 保持历史记录在合理范围内（最近10轮对话）
            if len(chat_history[chat_test_in.user_id]) > 21:  # system消息+20条对话消息
                chat_history[chat_test_in.user_id] = [
                    chat_history[chat_test_in.user_id][0]  # 保留system消息
                ] + chat_history[chat_test_in.user_id][
                    -20:
                ]  # 保留最近10轮对话

        # 返回流式响应
        return StreamingResponse(
            wrapped_generator(completion), media_type="text/event-stream"
        )

    except Exception as e:
        # 打印错误信息和文档参考链接
        print(f"错误信息：{e}")
        print(
            "请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code"
        )
        # 返回错误信息
        return {"error": str(e)}


class ChatPoemIn(BaseModel):
    content: str  # 诗句内容
    user_id: str  # 用户标识


@ai.post("/chat/poem/analysis", summary="AI赏析古诗")
# 使用@ai.post装饰器指定该函数处理的HTTP请求路径和方法，以及简要说明
async def chat_poem(chat_poem_in: ChatPoemIn):
    """
    接收一个诗句分析请求，使用指定的AI模型进行分析并返回分析结果。

    :param:
    - chat_poem_in: ChatPoemIn类型，包含需要分析的诗句内容。

    :return:
    - 成功时返回一个字典，包含分析结果和状态信息。
    - 失败时返回一个包含错误信息的字典。
    """
    try:
        # 初始化OpenAI客户端，使用特定的API密钥和基础URL
        client = OpenAI(
            api_key=ailiyun_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        # 使用AI模型进行诗句分析
        completion = client.chat.completions.create(
            model="qwen-plus",  # 指定使用的AI模型
            messages=[  # 构建对话消息列表，包括系统角色和用户角色的消息
                {"role": "system", "content": "你是一个大诗人，请赏析我给你发的诗句"},
                {"role": "user", "content": chat_poem_in.content},
            ],
            stream=False,  # 关闭流式输出，即一次性返回完整的结果
        )

        # 直接返回完整的回复内容
        return {
            "msg": "ok",
            "data": {
                "content": completion.choices[0].message.content,  # 提取AI回复的内容
            },
            "code": 200,  # HTTP状态码200，表示请求成功
        }

    except Exception as e:
        # 捕获异常，打印错误信息和错误代码文档链接
        print(f"错误信息：{e}")
        print(
            "请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code"
        )
        # 返回错误信息
        return {"error": str(e)}


# 用于存储诗歌对话的历史记录
poem_dialogue_history = {}


class ChatPoemDialogueIn(BaseModel):
    content: str  # 诗句内容
    name: str  # 诗名
    author: str  # 作者
    background: str  # 背景
    user_id: str  # 用户标识
    user_content: str  # 用户输入的内容


# 根据发送过来的诗句、作者，背景进行ai对话
@ai.post("/chat/poem/dialogue", summary="根据发送过来的诗句、作者，背景进行ai对话")
async def chat_poem_dialogue(chat_poem_dialogue_in: ChatPoemDialogueIn):
    """
    AI对话 endpoint，根据客户端发送的诗句、作者和背景信息进行对话。

    参数:
    - chat_poem_dialogue_in: ChatPoemDialogueIn 类型，包含用户ID、诗句名称、内容、作者、背景和用户内容。

    返回:
    - 成功时返回包含AI回复内容的字典。
    - 出错时返回错误信息。
    """
    try:
        # 初始化OpenAI客户端
        client = OpenAI(
            api_key=ailiyun_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        # 获取用户当前的对话历史
        current_history = poem_dialogue_history.get(chat_poem_dialogue_in.user_id, [])

        # 检查是否需要重置对话历史（新诗词或首次对话）
        should_reset = not current_history or (  # 首次对话
            len(current_history) > 0  # 检查是否是新诗词
            and f"《{chat_poem_dialogue_in.name}》诗：{chat_poem_dialogue_in.content}"
            not in current_history[0]["content"]
        )

        if should_reset:
            # 重置对话历史
            poem_dialogue_history[chat_poem_dialogue_in.user_id] = [
                {
                    "role": "system",
                    "content": f"假如你是诗人{chat_poem_dialogue_in.author}，你此时的背景是{chat_poem_dialogue_in.background}，你写下了这首《{chat_poem_dialogue_in.name}》诗：{chat_poem_dialogue_in.content}，我有一些问题想要问你，回答可以简练一点，不要使用markdown格式",
                }
            ]

        # 添加用户的新消息到历史记录
        poem_dialogue_history[chat_poem_dialogue_in.user_id].append(
            {"role": "user", "content": chat_poem_dialogue_in.user_content}
        )

        # 创建对话完成请求
        completion = client.chat.completions.create(
            model="qwen-plus",
            messages=poem_dialogue_history[chat_poem_dialogue_in.user_id],
            stream=False,  # 关闭流式输出
        )

        # 获取AI的回复
        assistant_response = completion.choices[0].message

        # 将助手的回复添加到历史记录
        poem_dialogue_history[chat_poem_dialogue_in.user_id].append(
            {"role": "assistant", "content": assistant_response.content}
        )

        # 保持历史记录在合理范围内（最近10轮对话）
        if len(poem_dialogue_history[chat_poem_dialogue_in.user_id]) > 21:
            poem_dialogue_history[chat_poem_dialogue_in.user_id] = [
                poem_dialogue_history[chat_poem_dialogue_in.user_id][0]
            ] + poem_dialogue_history[chat_poem_dialogue_in.user_id][-20:]

        # 返回成功响应
        return {
            "msg": "ok",
            "data": {
                "content": assistant_response.content,
            },
            "code": 200,
        }

    except Exception as e:
        # 打印错误信息并返回错误响应
        print(f"错误信息：{e}")
        print(
            "请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code"
        )
        return {"error": str(e)}


# 诗句补全的输入类型检测
class ChatPoemCompletionIn(BaseModel):
    content: str  # 诗句内容
    user_id: str  # 用户标识
    type: str  # 诗句类型
    name: str  # 诗句标题
    keywords: str  # 关键词


# 根据用书输入的诗句进行补全
@ai.post("/chat/poem/completion", summary="根据用书输入的诗句进行补全")
async def chat_poem_completion(chat_poem_completion_in: ChatPoemCompletionIn):
    content = chat_poem_completion_in.content
    name = chat_poem_completion_in.name
    keywords = chat_poem_completion_in.keywords
    type = chat_poem_completion_in.type

    try:
        # 初始化OpenAI客户端
        client = OpenAI(
            api_key=ailiyun_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        completion = client.chat.completions.create(
            model="qwen-plus",
            messages=[
                {
                    "role": "system",
                    "content": "假如你现在是一个会写古诗的大师，请根据用户输入的诗句进行补全，同时不要使用markdown格式，同时你会严格按照诗的格式，每句诗的最后一个字后面要加上标点符号来加一区分",
                },
                {
                    "role": "user",
                    "content": f"这首诗的名字是{name}，这首诗的关键词是{keywords}，这首诗的类型是{type}，目前已经写的内容是{content}，请根据这些信息对于诗句进行补全，只需要补全诗句内容，不能添加任何其他内容，请严格按照诗的格式",
                },
            ],
        )

        return {
            "msg": "ok",
            "data": {"content": completion.choices[0].message.content},
            "code": 200,
        }
    except Exception as e:
        print(f"错误信息：{e}")
        print(
            "请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code"
        )
        return {"error": str(e)}


# 根据用书输入的诗句提供建议
@ai.post("/chat/poem/advice", summary="根据用书输入的诗句提供建议")
async def chat_poem_advice(chat_poem_advice_in: ChatPoemCompletionIn):
    """
    异步处理用户输入的诗句，提供个性化的写作建议。
    
    参数:
    - chat_poem_advice_in: 包含用户输入诗句信息的请求体，包括内容、诗名、关键词和类型。
    
    返回:
    - 成功时返回一个包含建议内容的字典，包括消息状态、数据内容和状态码。
    - 失败时返回一个包含错误信息的字典。
    """
    # 提取用户输入的诗句内容、诗名、关键词和类型
    content = chat_poem_advice_in.content
    name = chat_poem_advice_in.name
    keywords = chat_poem_advice_in.keywords
    type = chat_poem_advice_in.type

    try:
        # 初始化OpenAI客户端，使用特定的API密钥和基础URL
        client = OpenAI(
            api_key=ailiyun_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        # 创建聊天完成请求，指定模型和对话消息
        completion = client.chat.completions.create(
            model="qwen-plus",
            messages=[
                {
                    "role": "system",
                    "content": "假如你现在是一个会写古诗的大师，请根据用户输入的诗句提供建议，同时不要使用markdown格式，同时你会严格按照诗的格式",
                },
                {
                    "role": "user",
                    "content": f"这首诗的名字是{name}，这首诗的关键词是{keywords}，这首诗的类型是{type}，目前已经写的内容是{content}，请根据这些信息对于诗句写作提供建议",
                },
            ],
        )
        # 返回成功响应，包含建议内容
        return {
            "msg": "ok",
            "data": {"content": completion.choices[0].message.content},
            "code": 200,
        }
    except Exception as e:
        # 捕获异常，打印错误信息和文档参考链接
        print(f"错误信息：{e}")
        print(
            "请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code"
        )
        # 返回错误响应，包含错误信息
        return {"error": str(e)}
